package com.springboot.uav.util;

import net.sourceforge.tess4j.ITesseract;
import net.sourceforge.tess4j.Tesseract;
import net.sourceforge.tess4j.TesseractException;
import org.springframework.stereotype.Component;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.File;
import java.io.IOException;
import java.net.URL;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.regex.Pattern;

@Component
public class OcrUtil {

    // Tesseract训练数据路径，需根据实际情况修改
    private static final String TESSDATA_PATH = "tessdata";

    // 用于将彩色图像转换为灰度图时，红色通道的权重
    private static final double RED_WEIGHT = 0.2126;
    // 用于将彩色图像转换为灰度图时，绿色通道的权重
    private static final double GREEN_WEIGHT = 0.7152;
    // 用于将彩色图像转换为灰度图时，蓝色通道的权重
    private static final double BLUE_WEIGHT = 0.0722;

    // 缓存Tesseract实例，避免每次调用都重新创建，提高性能
    private static ITesseract tesseract;

    static {
        try {
            tesseract = new Tesseract();
            tesseract.setDatapath(TESSDATA_PATH);
        } catch (Exception e) {
            // 初始化Tesseract实例失败时，打印异常信息，可根据实际需求更好地处理
            e.printStackTrace();
        }
    }

    /**
     * 从图片的URL中提取取件码信息
     *
     * @param imageUrl 图片的URL地址
     * @return 提取到的取件码列表，如果未提取到则返回空列表
     * @throws TesseractException 如果Tesseract执行OCR操作时出现异常
     * @throws IOException        如果读取图片时出现异常
     */
    public static List<String> extractCodeFromImageUrl(String imageUrl) throws TesseractException, IOException {
        // 从URL读取图片并转换为BufferedImage
        BufferedImage image = readImageFromUrl(imageUrl);

        // 对图片进行预处理，包括转换为灰度图和降噪等操作
        BufferedImage processedImage = preprocessImage(image);

        // 使用缓存的Tesseract实例进行OCR操作，获取识别出的文本
        String recognizedText = tesseract.doOCR(processedImage);

        // 从识别出的文本中提取取件码
        return extractCodesFromText(recognizedText);
    }

    /**
     * 从给定的URL读取图片并转换为BufferedImage对象
     *
     * @param imageUrl 图片的URL地址
     * @return 读取到的BufferedImage对象
     * @throws IOException 如果读取图片时出现网络连接等异常
     */
    private static BufferedImage readImageFromUrl(String imageUrl) throws IOException {
        // 设置最大重试次数，当读取图片失败时尝试重新读取
        int maxRetries = 3;
        for (int i = 0; i < maxRetries; i++) {
            try {
                return ImageIO.read(new URL(imageUrl));
            } catch (IOException e) {
                if (i == maxRetries - 1) {
                    throw e;
                }
                // 打印读取图片失败信息，可根据实际需求添加更详细的日志记录
                System.out.println("读取图片失败，正在重试...");
                try {
                    Thread.sleep(1000);
                } catch (InterruptedException ex) {
                    Thread.currentThread().interrupt();
                }
            }
        }
        // 正常情况下不会执行到这里，只是为了满足编译器要求
        return null;
    }

    /**
     * 对图片进行预处理操作，包括转换为灰度图和降噪处理
     *
     * @param image 要处理的原始图片的BufferedImage对象
     * @return 处理后的图片的BufferedImage对象
     */
    private static BufferedImage preprocessImage(BufferedImage image) {
        // 先将图片转换为灰度图
        image = convertToGrayScale(image);

        // 再对灰度图进行降噪处理，这里使用中值滤波算法
        image = medianFilter(image);

        return image;
    }

    /**
     * 将彩色图片转换为灰度图的方法
     *
     * @param image 要转换的彩色图片的BufferedImage对象
     * @return 转换后的灰度图的BufferedImage对象
     */
    private static BufferedImage convertToGrayScale(BufferedImage image) {
        int width = image.getWidth();
        int height = image.getHeight();

        // 创建一个新的灰度图BufferedImage对象，类型为TYPE_BYTE_GRAY
        BufferedImage grayScaleImage = new BufferedImage(width, height, BufferedImage.TYPE_BYTE_GRAY);

        for (int y = 0; y < height; y++) {
            for (int x = 0; x < width; x++) {
                int argb = image.getRGB(x, y);
                int alpha = (argb >> 24) & 0xff;
                int red = (argb >> 16) & 0xff;
                int green = (argb >> 8) & 0xff;
                int blue = argb & 0xff;

                // 根据设定的权重计算灰度值
                int gray = (int) (RED_WEIGHT * red + GREEN_WEIGHT * green + BLUE_WEIGHT * blue);

                // 设置灰度图中对应像素的颜色值
                grayScaleImage.setRGB(x, y, (alpha << 24) | (gray << 16) | (gray << 8) | gray);
            }
        }

        return grayScaleImage;
    }

    /**
     * 对图片进行中值滤波降噪处理的方法
     *
     * @param image 要降噪处理的图片的BufferedImage对象
     * @return 降噪处理后的图片的BufferedImage对象
     */
    private static BufferedImage medianFilter(BufferedImage image) {
        int width = image.getWidth();
        int height = image.getHeight();

        // 创建一个新的BufferedImage对象用于存储降噪后的图片，类型为TYPE_INT_ARGB
        BufferedImage filteredImage = new BufferedImage(width, height, BufferedImage.TYPE_INT_ARGB);

        // 中值滤波的窗口大小，这里设置为3x3
        int size = 3;
        int halfSize = size / 2;

        for (int y = halfSize; y < height - halfSize; y++) {
            for (int x = halfSize; x < width - halfSize; x++) {
                // 用于存储窗口内每个像素的红色、绿色、蓝色通道值
                int[] red = new int[size * size];
                int[] green = new int[size * size];
                int[] blue = new int[size * size];
                int k = 0;

                // 遍历窗口内的每个像素，获取其颜色通道值
                for (int j = -halfSize; j <= halfSize; j++) {
                    for (int i = -halfSize; i <= halfSize; i++) {
                        int argb = image.getRGB(x + i, y + j);
                        red[k] = (argb >> 16) & 0xff;
                        green[k] = (argb >> 8) & 0xff;
                        blue[k] = argb & 0xff;
                        k++;
                    }
                }

                // 对每个颜色通道的值进行排序
                Arrays.sort(red);
                Arrays.sort(green);
                Arrays.sort(blue);

                // 取排序后的中间值作为中值滤波后的颜色通道值
                int medianRed = red[size * size / 2];
                int medianGreen = green[size * size / 2];
                int medianBlue = blue[size * size / 2];

                // 构建降噪后像素的颜色值
                int medianArgb = (0xff << 24) | (medianRed << 16) | (medianGreen << 8) | medianBlue;

                // 设置降噪后图片中对应像素的颜色值
                filteredImage.setRGB(x, y, medianArgb);
            }
        }

        return filteredImage;
    }

    /**
     * 从识别出的文本中提取符合取件码格式的字符串
     *
     * @param text 识别出的文本内容
     * @return 提取到的取件码列表
     */
    private static List<String> extractCodesFromText(String text) {
        // 定义取件码的可能格式正则表达式，这里示例为数字和横线组成的格式，可根据实际情况修改
        Pattern codePattern = Pattern.compile("\\d+-\\d+-\\d+");
        List<String> codes = new ArrayList<>();

        for (String line : text.split("\n")) {
            java.util.regex.Matcher matcher = codePattern.matcher(line);
            while (matcher.find()) {
                codes.add(matcher.group());
            }
        }

        return codes;
    }
}
